#!/usr/bin/env python3
import rospy
import threading
import torch
from func.f_process import Process
from std_msgs.msg import Float32MultiArray
from sensor_msgs.msg import Image
from uav_event.msg import UInt32Mat
from rich import print


publisherInitialized = False
nodePub = None


subscriberInitialized = False
nodeSub = None


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
fsnn = Process(device = device)


sumTime = 0.0
counter = -20
minDuration = torch.inf
maxDuration = 0


def nodeSubCallback(msg: UInt32Mat) -> None:
    global fsnn, publisherInitialized, nodePub, device, sumTime, counter, minDuration, maxDuration

    eventData = torch.tensor(msg.data, dtype = torch.long, device = device)
    if eventData.shape[0]:
        eventSize = len(msg.column)
        eventData = eventData.reshape(eventData.shape[0] // eventSize, eventSize) # [sec, nsec, p, y, x]

    res = fsnn.execute(eventData)

    pubRes = Float32MultiArray()
    pubRes.data = res.flatten().cpu().numpy().tolist()
    if pubRes.data[0]:
        if publisherInitialized:
            print("[green bold]Succesfully send force as message.[/green bold]")
            nodePub.publish(pubRes)
        else:
            print("[purple bold]Force generated, but blocked.[/purple bold]")
        print(pubRes.data)
        print("[purple bold]Refractory time, no input will be accepted.[/purple bold]")

    if counter > 0:
        duration = sum(fsnn.processTimes)
        sumTime += duration
        if minDuration > duration:
            minDuration = duration
        if maxDuration < duration:
            maxDuration = duration
        print("SNN: min %.6f ms, max %.6f ms, avr %.6f ms" % (1000 * minDuration, 1000 * maxDuration, 1000 * sumTime / counter))
    counter += 1


def subThreadFunc() -> None:
    global subscriberInitialized, nodeSub
    nodeSub = rospy.Subscriber("/uav_event/ego_motion_comp_res", UInt32Mat, nodeSubCallback)
    subscriberInitialized = True
    rospy.spin()


def pubThreadFunc() -> None:
    global publisherInitialized, nodePub
    nodePub = rospy.Publisher("/uav_event/execution_res", Float32MultiArray, queue_size = 1)
    publisherInitialized = True


if __name__ == "__main__":
    rospy.init_node("process_snn_controller", anonymous = True)

    subT = threading.Thread(target = subThreadFunc)
    pubT = threading.Thread(target = pubThreadFunc)

    subT.start()
    pubT.start()

    subT.join()
    pubT.join()
